﻿using System;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Rendering;

namespace Mars.GpuInstance
{
    public abstract class GPURenderBatch : IDisposable, IRecyclable
    {
        protected static Matrix4x4[] matrices = new Matrix4x4[InstancerConstants.COUNT_PER_DRAW];
        protected static Vector4[] campColorDatas = new Vector4[InstancerConstants.COUNT_PER_DRAW];
        protected static Bounds defaultBound = new Bounds(Vector3.zero, Vector3.one);

        public string meshPath;
        public Mesh renderMesh;
        public Material renderMat;
        public float2 drawRange;
        protected MaterialPropertyBlock mpb;
        protected NativeArray<float4x4> output;
        protected NativeArray<float4> campData;
        protected NativeArray<int> outputCount;
        protected NativeList<int> outputIndex;
        protected NativeArray<float4> boundVerts;
        protected Matrix4x4 transformOffset;
        protected bool hasPoseOffset = false;
        protected Matrix4x4 bindPoseOffset;
        protected Bounds meshBounds;
        
        protected bool waitingSubmitVisibilityBuffer = false;
        
        //包围盒扩大的值
        protected float3 boxExtends = float3.zero;
        
        protected int totalBoneCount;

        public NativeArray<float> minMaxDis = new(2, Allocator.Persistent);
        public int lodLevel;
        
        public bool isSkinBatch;
        public bool nonProportionalScaling;     // 非等比缩放，针对这种情况需要开启一个 keyword，因此单独作为一个 batch
        
        protected int errorCounter;
        
        public abstract void CopyAnimationData(in NativeList<float4> animatorControllerData,
            in NativeList<float4> animationData,
            NativeList<JobHandle> handles);

        public abstract JobHandle CalculateAnim(int frameRate, int textureSizeX, int totalBoneCount);

        public abstract void UpdateAnimationBakeBuffer();

        public void Init(Mesh _mesh, Material _mat, MaterialPropertyBlock _mpb, float2 _drawRange, float3 _boxExtends)
        {
            meshBounds = _mesh.bounds;
            renderMesh = _mesh;
  
            renderMat = _mat;
            drawRange = _drawRange;
            mpb = _mpb;
            boxExtends = _boxExtends;
            nonProportionalScaling = false;
            meshPath = "";
            
            isSkinBatch = this is GPUSkinRenderBatch;

            hasPoseOffset = false;
            
            errorCounter = 0;

#if UNITY_EDITOR || UNITY_STANDALONE_WIN
            if (!_mat.enableInstancing)
            {
                Debug.LogError($"[GPURenderBatch] 材质 {_mat.name} 不支持gpu instancing. 找美术在材质上勾上Enable GPU Instancing");
            }          
#endif

        }

        public void Init(string _meshPath, Material _mat, MaterialPropertyBlock _mpb, float2 _drawRange, float3 _boxExtends)
        {
            renderMesh = null;
            renderMat = _mat;
            drawRange = _drawRange;
            mpb = _mpb;
            boxExtends = _boxExtends;
            nonProportionalScaling = false;
            meshPath = _meshPath;
            meshBounds = defaultBound;

            isSkinBatch = this is GPUSkinRenderBatch;

            hasPoseOffset = false;
            
            errorCounter = 0;

#if UNITY_EDITOR || UNITY_STANDALONE_WIN
            if (!_mat.enableInstancing)
            {
                Debug.LogError($"[GPURenderBatch] 材质 {_mat.name} 不支持gpu instancing. 找美术在材质上勾上Enable GPU Instancing");
            }          
#endif

        }

        public void Init(GPURenderBatch other, bool nonPropScaling)
        {
            meshBounds = other.meshBounds;
            renderMesh = other.renderMesh;
            renderMat = new Material(other.renderMat);
            drawRange = other.drawRange;
            mpb = other.mpb;
            boxExtends = other.boxExtends;
            meshPath = other.meshPath;
            
            transformOffset = other.transformOffset;
            totalBoneCount = other.totalBoneCount;
            lodLevel = other.lodLevel;
            isSkinBatch = other.isSkinBatch;

            UpdateNonPropScalingKeyword(nonPropScaling);
            errorCounter = 0;
        }

        public virtual void SubmitToGPU()
        {
            if (waitingSubmitVisibilityBuffer)
            {
                SubmitVisibilityBuffer();
            }
        }
        
        public void VisibilityCalculation(in GPUAnimatorInstancerRuntimeData runtimeData, in NativeArray<float4> frustumPlanes, in float3 cameraPos, NativeList<JobHandle> handles)
        {
            if (!output.IsCreated)
            {
                output = new NativeArray<float4x4>(runtimeData.input.Length, Allocator.Persistent);
            }
            else if (output.IsCreated && output.Length < runtimeData.input.Length)
            {
                output.Dispose();
                output = new NativeArray<float4x4>(runtimeData.input.Length, Allocator.Persistent);
            }
            
            if (!campData.IsCreated)
            {
                campData = new NativeArray<float4>(runtimeData.input.Length, Allocator.Persistent);
            }
            else if (campData.IsCreated && campData.Length < runtimeData.input.Length)
            {
                campData.Dispose();
                campData = new NativeArray<float4>(runtimeData.input.Length, Allocator.Persistent);
            }

            waitingSubmitVisibilityBuffer = true;
            
            outputIndex.Clear();
            
            var job = new SimpleAnimatorFrustumCullingJob();
            job.input = runtimeData.input;
            job.inputCampData = runtimeData.campColorData;
            job.outputIndex = outputIndex;
            job.boxCenter = meshBounds.center;
            job.boxExtents = (float3)(meshBounds.extents) + boxExtends;
            job.cameraPlanes = frustumPlanes;
            job.drawRange = drawRange;
            job.cameraPos = cameraPos;
            job.output = output;
            job.boundVerts = boundVerts;
            job.campData = campData;
            job.transformOffset = transformOffset;

            job.isSkinBatch = isSkinBatch;
            job.batchNonPropScaling = nonProportionalScaling;
            job.inputNonPropScaling = runtimeData.inputNonPropScaling;
            
            handles.Add(job.Schedule());
        }

        public int GetInstanceCount()
        {
            if (outputIndex.IsCreated)
            {
                return outputIndex.Length;
            }

            return 0;
        }

        public void UpdateBoneOffset(Matrix4x4 _bindPoseOffset)
        {
            hasPoseOffset = true;
            bindPoseOffset = _bindPoseOffset;
            renderMat.EnableKeyword(InstancerConstants.KEYWORD_GPU_BINDPOSEOFFSET);
            //mpb.SetMatrix(InstancerConstants.BINDPOSE_OFFSET, bindPoseOffset);
        }

        public void UpdateTransformOffset(Matrix4x4 offset)
        {
            this.transformOffset = offset;
            //mpb.SetMatrix(InstancerConstants.TRANSFORM_OFFSET, offset);
        }

        public void UpdateNonPropScalingKeyword(bool nonPropScaling)
        {
            nonProportionalScaling = nonPropScaling;
            renderMat.EnableKeyword(InstancerConstants.KEYWORD_ANI_NON_PROP_SCALING);
        }
        
        protected void SubmitVisibilityBuffer()
        {
            waitingSubmitVisibilityBuffer = false;
        }

        public virtual void Dispose()
        {
            if (campData.IsCreated)
            {
                campData.Dispose();
            }

            if (output.IsCreated)
            {
                output.Dispose();
            }

            if (minMaxDis.IsCreated)
            {
                minMaxDis.Dispose();
            }
            
            if (outputCount.IsCreated)
            {
                outputCount.Dispose();
            }

            if (outputIndex.IsCreated)
            {
                outputIndex.Dispose();
            }
            
            if (boundVerts.IsCreated)
            {
                boundVerts.Dispose();
            }

            renderMat = null;
            renderMesh = null;
            mpb = null;
            hasPoseOffset = false;
        }

        #region Draw

        public abstract void Draw(Bounds drawBounds, ShadowLod animatorShadowLod);
        public abstract void DrawDebug(float sightDistance);
        public abstract void DrawShadows(Bounds drawBounds);
        #endregion
        
        #region 用DrawMeshInstanced绘制

        protected void DrawMeshInstanced(Matrix4x4[] matrices, int count, ShadowCastingMode shadowCastingMode)
        {
            if (!renderMat || !renderMesh)
            {
                if (errorCounter == 0)
                {
                    Debug.LogError($"[{nameof(GPURenderBatch)}] DrawMeshInstanced failed, renderMat = {renderMat}, renderMesh = {renderMesh}");
                    errorCounter++;
                }
                return;
            }
            var subMeshCount = renderMesh.subMeshCount;
            for (int j = 0; j < subMeshCount; j++)
            {
                Graphics.DrawMeshInstanced(
                    renderMesh, 
                    j, 
                    renderMat, 
                    matrices, 
                    count, 
                    mpb, 
                    shadowCastingMode);
            }
        }

        #endregion

        public virtual void OnRecycle()
        {
            if (outputIndex.IsCreated)
            {
                outputIndex.Clear();
            }

            renderMat = null;
            renderMesh = null;
            mpb = null;

            hasPoseOffset = false;
        }
    }
}
